from typing import List
from app.models.knowledge_base import Document
from app.services.embedding_service import EmbeddingService
from sqlalchemy import text
from app import db
import numpy as np
import logging

logger = logging.getLogger(__name__)

class VectorStore:
    @staticmethod
    def search(query_vector: List[float], kb_id: int, max_results: int = 5, similarity_threshold: float = 0.5) -> List[Document]:
        """
        搜索最相似的文档
        """
        try:
            logger.info(f"Searching document for kb_id={kb_id}, max_results={max_results}")
            
            # 将查询向量转换为字符串
            vector_str = f"[{','.join(map(str, query_vector))}]"
            
            # 构建查询SQL，使用余弦相似度
            sql = text(f"""
                SELECT id, content, 
                    1 - (embedding <#> '{vector_str}'::vector) AS cosine_similarity
                FROM public.document
                WHERE knowledge_base_id = :kb_id
                    AND embedding IS NOT NULL
                ORDER BY cosine_similarity DESC
                LIMIT :limit
            """)
            
            # 执行查询
            result = db.session.execute(
                sql,
                {
                    'kb_id': kb_id,
                    'limit': max_results
                }
            )
            
            # 获取文档ID列表和相似度
            doc_similarities = [(row.id, row.cosine_similarity) for row in result]
            
            if not doc_similarities:
                logger.info("No documents found")
                return []
            
            # 获取文档详情
            doc_ids = [doc_id for doc_id, _ in doc_similarities]
            docs = Document.query.filter(Document.id.in_(doc_ids)).all()
            
            # 按相似度过滤和排序
            filtered_docs = []
            for doc in docs:
                similarity = next(sim for id_, sim in doc_similarities if id_ == doc.id)
                if similarity >= similarity_threshold:
                    doc.similarity = similarity  # 添加相似度属性
                    filtered_docs.append(doc)
            
            # 按相似度降序排序
            filtered_docs.sort(key=lambda x: x.similarity, reverse=True)
            
            logger.info(f"Found {len(filtered_docs)} documents above similarity threshold")
            return filtered_docs
            
        except Exception as e:
            logger.error(f"Error searching documents: {str(e)}")
            raise

    @staticmethod
    def add_document(document: Document) -> bool:
        """
        添加文档到向量存储
        
        Args:
            document: 文档对象
            
        Returns:
            bool: 是否添加成功
        """
        try:
            db.session.add(document)
            db.session.commit()
            logger.info(f"Added document {document.id} to vector store")
            return True
        except Exception as e:
            logger.error(f"Error adding document: {str(e)}")
            db.session.rollback()
            raise

    @staticmethod
    def delete_document(document_id: int) -> bool:
        """
        从向量存储中删除文档
        
        Args:
            document_id: 文档ID
            
        Returns:
            bool: 是否删除成功
        """
        try:
            doc = Document.query.get(document_id)
            if doc:
                db.session.delete(doc)
                db.session.commit()
                logger.info(f"Deleted document {document_id} from vector store")
                return True
            return False
        except Exception as e:
            logger.error(f"Error deleting document: {str(e)}")
            db.session.rollback()
            raise 